// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

namespace kumo::pollux::functions {
namespace {

class MapMatchTest : public functions::test::FunctionBaseTest {
 protected:
  void match(
      const std::string& functionName,
      const VectorPtr& input,
      const std::string& lambda,
      const std::vector<std::optional<bool>>& expected) {
    const std::string expr =
        fmt::format("{}(c0, x -> ({}))", functionName, lambda);

    SCOPED_TRACE(expr);
    auto result = evaluate(expr, make_row_vector({input}));
    pollux::test::assertEqualVectors(
        make_nullable_flat_vector<bool>(expected), result);
  }

  void anyKeysMatch(
      const VectorPtr& input,
      const std::string& lambda,
      const std::vector<std::optional<bool>>& expected) {
    match("any_keys_match", input, lambda, expected);
  }

  void allKeysMatch(
      const VectorPtr& input,
      const std::string& lambda,
      const std::vector<std::optional<bool>>& expected) {
    match("all_keys_match", input, lambda, expected);
  }

  void noKeysMatch(
      const VectorPtr& input,
      const std::string& lambda,
      const std::vector<std::optional<bool>>& expected) {
    match("no_keys_match", input, lambda, expected);
  }

  void anyValuesMatch(
      const VectorPtr& input,
      const std::string& lambda,
      const std::vector<std::optional<bool>>& expected) {
    match("any_values_match", input, lambda, expected);
  }

  void noValuesMatch(
      const VectorPtr& input,
      const std::string& lambda,
      const std::vector<std::optional<bool>>& expected) {
    match("no_values_match", input, lambda, expected);
  }
};

TEST_F(MapMatchTest, anyKeysMatch) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 10, -2: 20}",
  });

  anyKeysMatch(data, "x > 0", {true, false});
  anyKeysMatch(data, "x > 2", {true, false});
  anyKeysMatch(data, "x > 3", {false, false});
  anyKeysMatch(data, "x <= 0", {false, true});
  anyKeysMatch(data, "x <= -2", {false, true});
  anyKeysMatch(data, "x <= -5", {false, false});
}

TEST_F(MapMatchTest, anyKeysMatchNull) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 10, -2: 20}",
  });

  anyKeysMatch(data, "if(x = 2, null::boolean, false)", {std::nullopt, false});
  anyKeysMatch(data, "if(x = 2, null::boolean, true)", {true, true});
}

TEST_F(MapMatchTest, allKeysMatch) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 10, -2: 20}",
  });

  allKeysMatch(data, "x > 0", {true, false});
  allKeysMatch(data, "x > 1", {false, false});
  allKeysMatch(data, "x > 10", {false, false});
  allKeysMatch(data, "x <= 0", {false, true});
  allKeysMatch(data, "x <= 3", {true, true});
  allKeysMatch(data, "x <= 100", {true, true});
}

TEST_F(MapMatchTest, allKeysMatchNull) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 10, -2: 20}",
  });

  allKeysMatch(data, "if(x = 2, null::boolean, false)", {false, false});
  allKeysMatch(data, "if(x = 2, null::boolean, true)", {std::nullopt, true});
}

TEST_F(MapMatchTest, noKeysMatch) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 10, -2: 20}",
  });

  noKeysMatch(data, "x > 0", {false, true});
  noKeysMatch(data, "x = 2", {false, true});
  noKeysMatch(data, "x = 22", {true, true});
  noKeysMatch(data, "x = -2", {true, false});
  noKeysMatch(data, "x > 2 OR x < -1", {false, false});
}

TEST_F(MapMatchTest, noKeysMatchNull) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 10, -2: 20}",
  });

  noKeysMatch(data, "if(x = 2, null::boolean, false)", {std::nullopt, true});
  noKeysMatch(data, "if(x = 2, null::boolean, true)", {false, false});
}

TEST_F(MapMatchTest, anyValuesMatch) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 11, -2: 22}",
  });

  anyValuesMatch(data, "x = 10", {true, false});
  anyValuesMatch(data, "x = 22", {false, true});
  anyValuesMatch(data, "x < 15", {true, true});
  anyValuesMatch(data, "x < 0", {false, false});
  anyValuesMatch(data, "x IN (20, 11)", {true, true});
}

TEST_F(MapMatchTest, anyValuesMatchNull) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 11, -2: 22}",
  });

  anyValuesMatch(
      data, "if(x = 20, null::boolean, false)", {std::nullopt, false});
  anyValuesMatch(data, "if(x = 20, null::boolean, true)", {true, true});
}

TEST_F(MapMatchTest, noValuesMatch) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 11, -2: 22}",
  });

  noValuesMatch(data, "x = 7", {true, true});
  noValuesMatch(data, "x > 15", {false, false});
  noValuesMatch(data, "x > 25", {false, true});
  noValuesMatch(data, "x % 11 = 0", {true, false});
}

TEST_F(MapMatchTest, noValuesMatchNull) {
  auto data = make_map_vector_from_json<int32_t, int64_t>({
      "{1: 10, 2: 20, 3: 30}",
      "{-1: 11, -2: 22}",
  });

  noValuesMatch(data, "if(x = 20, null::boolean, false)", {std::nullopt, true});
  noValuesMatch(data, "if(x = 20, null::boolean, true)", {false, false});
}

} // namespace
} // namespace kumo::pollux::functions
